export const hitl = (options: HITLPluginOptions = {}): HITLController => {
  let mode: HITLMode = options.mode ?? 'ask'
  const callAllows = new Set<string>()
  const pending = new Map<string, HITLPendingRequest>()
  const pendingConversationAllows = new Set<string>()
  const resumeDecisions = new Map<string, HITLResumeDecision>()
  const rejectedKeys = new Map<string, string>()
  const runAllows = new Set<string>()

  const resolveDecision = async (context: HITLPolicyContext): Promise<HITLPolicyDecision> => {
    if (mode === 'off' || mode === 'allow')
      return { type: 'allow' }
    if (mode === 'deny')
      return { type: 'deny' }

    return await options.policy?.(context) ?? { type: 'ask' }
  }

  const createControl = (privateState: PluginPrivateStateApi<HITLPrivateState>): ToolCallControl => ({
    preToolCall: async (context) => {
      const key = getToolKey(context.toolCall)

      if (mode === 'deny') {
        return {
          message: 'TOOL_HITL_REJECTED',
          type: 'reject',
        }
      }

      const rejectedMessage = rejectedKeys.get(key)
      if (rejectedMessage != null) {
        rejectedKeys.delete(key)
        return {
          message: rejectedMessage,
          type: 'reject',
        }
      }

      if (callAllows.has(key)) {
        callAllows.delete(key)
        return { type: 'continue' }
      }

      if (runAllows.has(key) || hasConversationAllow(privateState, key))
        return { type: 'continue' }

      const decision = await resolveDecision({
        ...context,
        key,
        metadata: metadata.get(context.tool),
      })

      if (decision.type === 'allow')
        return { type: 'continue' }

      if (decision.type === 'deny') {
        return {
          message: decision.message ?? 'TOOL_HITL_REJECTED',
          type: 'reject',
        }
      }

      const interruption: ToolInterruption = {
        data: {
          key,
          metadata: metadata.get(context.tool),
        },
        id: `hitl_${context.toolCall.id}`,
        reason: decision.message ?? 'Human review required.',
        toolCall: context.toolCall,
      }

      pending.set(interruption.id, {
        id: interruption.id,
        key,
        metadata: metadata.get(context.tool),
        reason: interruption.reason,
        toolCall: context.toolCall,
        toolName: context.tool.function.name,
      })

      return {
        interruption,
        type: 'interrupt',
      }
    },
  })

  const approve: HITLController['approve'] = (id, scope = options.scope ?? 'call') => {
    const request = pending.get(id)
    if (request == null)
      return false

    if (scope === 'call') {
      callAllows.add(request.key)
    }
    else if (scope === 'run') {
      runAllows.add(request.key)
    }
    else {
      pendingConversationAllows.add(request.key)
    }

    resumeDecisions.set(id, { id, scope, type: 'approved' })
    pending.delete(id)
    return true
  }

  const reject: HITLController['reject'] = (id, message = 'TOOL_HITL_REJECTED') => {
    const request = pending.get(id)
    if (request == null)
      return false

    rejectedKeys.set(request.key, message)
    resumeDecisions.set(id, {
      id,
      key: request.key,
      message,
      type: 'rejected',
    })
    return pending.delete(id)
  }

  return {
    approve,
    clear: () => {
      pending.clear()
      pendingConversationAllows.clear()
      callAllows.clear()
      rejectedKeys.clear()
      resumeDecisions.clear()
      runAllows.clear()
    },
    getDecisionForResume: id => resumeDecisions.get(id),
    pending: () => [...pending.values()],
    plugin: {
      name,
      onTurnDone: () => {
        runAllows.clear()
      },
      resolveToolCallControl: ({ privateState }) => {
        if (privateState == null)
          throw new Error('@apeira/plugin-hitl requires plugin private state support from @apeira/core.')

        if (pendingConversationAllows.size > 0) {
          for (const key of pendingConversationAllows) {
            setConversationAllow(privateState, key)
          }
          pendingConversationAllows.clear()
        }

        return createControl(privateState)
      },
      version,
    },
    reject,
    setMode: nextMode => mode = nextMode,
  }
}

const getToolKey = (toolCall: ToolCall) => {
  let args = toolCall.function.arguments ?? ''
  try {
    const parsed = JSON.parse(args)
    const canonicalize = (obj: any): any => {
      if (Array.isArray(obj))
        return obj.map(canonicalize)
      if (obj !== null && typeof obj === 'object') {
        return Object.keys(obj).sort().reduce((acc: any, key) => {
          acc[key] = canonicalize(obj[key])
          return acc
        }, {})
      }
      return obj
    }
    args = JSON.stringify(canonicalize(parsed))
  }
  catch {}
  return JSON.stringify([
    toolCall.function.name,
    args,
  ])
}